"""
作者：Leagolas
日期：2024年06月12日
"""
import pandas as pd
import numpy as np

def calculate_bollinger_bands(pd_data, factor, span=20):
    
    pd_data["平均值"] = np.nan
    pd_data["标准差"] = np.nan
    pd_data["上轨"] = np.nan
    pd_data["下轨"] = np.nan
    pd_data["决策"] = np.nan


    pd_data["平均值"] = pd_data[factor].rolling(window=span).mean()
    pd_data["标准差"] = pd_data[factor].rolling(window=span).std()
    pd_data["上轨"] = pd_data["平均值"] + 1 * pd_data["标准差"]
    pd_data["下轨"] = pd_data["平均值"] - 1 * pd_data["标准差"]


    pd_data["状态"] = np.where(pd_data[factor] >= pd_data["上轨"], 2, np.where(pd_data[factor] > pd_data["平均值"], 1, np.where(pd_data[factor] == pd_data["平均值"], 0 , np.where(pd_data[factor] > pd_data["下轨"], -1, -2))))
    pd_data.loc[0:span-2,"状态"] = np.nan

    return pd_data[['交易日',"状态"]]